// Copyright 2022 jeff.li. and/or its affiliates.
/*
 * Acknowledgement:
 * The structure of the expressions is inspired by Halide/TVM IR.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file tbir/ir/type_functor.h
 * \brief A way to defined arbitrary function signature with dispatch on types.
 */
#pragma once

#include <string>
#include <utility>
#include <vector>

#include <tbir/ir/adt.h>
#include <tbir/ir/base.h>
#include <tbir/runtime/container.h>
#include <tbir/runtime/functor.h>

namespace tbir::ir {

    using ::tbir::runtime::NodeFunctor;

    template<typename FType>
    class TypeFunctor;

// functions to be overridden.
#define TBIR_TYPE_FUNCTOR_DEFAULT \
  { return VisitTypeDefault_(op, std::forward<Args>(args)...); }

#define TBIR_TYPE_FUNCTOR_DISPATCH(OP)                                               \
  vtable.template set_dispatch<OP>([](const ObjectRef& n, TSelf* self, Args... args) {     \
    return self->VisitType_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
  });

    template<typename R, typename... Args>
    class TypeFunctor<R(const Type &n, Args...)> {
    private:
        using TSelf = TypeFunctor<R(const Type &n, Args...)>;
        using FType = NodeFunctor<R(const ObjectRef &n, TSelf *self, Args...)>;

    public:
        /*! \brief the result type of this functor */
        using result_type = R;

        /*! \brief virtual destructor */
        virtual ~TypeFunctor() {
        }

        /*!
         * \brief Same as call.
         * \param n The expression node.
         * \param args Additional arguments.
         * \return The result of the call
         */
        R operator()(const Type &n, Args... args) {
            return VisitType(n, std::forward<Args>(args)...);
        }

        /*!
         * \brief The functor call.
         * \param n The expression node.
         * \param args Additional arguments.
         * \return The result of the call
         */
        virtual R VisitType(const Type &n, Args... args) {
            MXCHECK(n.defined());
            static FType vtable = InitVTable();
            return vtable(n, this, std::forward<Args>(args)...);
        }

        // Functions that can be overriden by subclass
        virtual R VisitType_(const TypeVarNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const TypeConstraintNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const FuncTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const TupleTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const GlobalTypeVarNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const PrimTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const PointerTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const ObjectTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const UnicodeTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const StringTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const ListTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const DictTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const SetTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const IteratorTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const ExceptionTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const FileTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const NDArrayTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const ClassTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const UserDataTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const OpaqueObjectTypeNode *op,
                             Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitType_(const RefTypeNode *op, Args... args) TBIR_TYPE_FUNCTOR_DEFAULT;

        virtual R VisitTypeDefault_(const Object *op, Args...) {
            MXLOG(FATAL) << "Do not have a default for " << op->GetTypeKey();
            throw;  // unreachable, written to stop compiler warning
        }

    private:
        // initialize the vtable.
        static FType InitVTable() {
            FType vtable;
            // Set dispatch
            TBIR_TYPE_FUNCTOR_DISPATCH(TypeVarNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(TypeConstraintNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(FuncTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(TupleTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(GlobalTypeVarNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(PrimTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(PointerTypeNode);

            TBIR_TYPE_FUNCTOR_DISPATCH(ObjectTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(UnicodeTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(StringTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(ListTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(DictTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(SetTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(IteratorTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(ExceptionTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(FileTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(NDArrayTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(ClassTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(UserDataTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(OpaqueObjectTypeNode);
            TBIR_TYPE_FUNCTOR_DISPATCH(RefTypeNode);
            return vtable;
        }
    };

#undef TBIR_TYPE_FUNCTOR_DISPATCH

    /*!
     * \brief A type visitor that recursively visit types.
     */
    class TURBO_EXPORT TypeVisitor : public TypeFunctor<void(const Type &n)> {
    public:
        void VisitType_(const TypeVarNode *op) override;

        void VisitType_(const FuncTypeNode *op) override;

        void VisitType_(const TupleTypeNode *op) override;

        void VisitType_(const GlobalTypeVarNode *op) override;

        void VisitType_(const PrimTypeNode *op) override;

        void VisitType_(const PointerTypeNode *op) override;

        void VisitType_(const ObjectTypeNode *op) override;

        void VisitType_(const UnicodeTypeNode *op) override;

        void VisitType_(const StringTypeNode *op) override;

        void VisitType_(const ListTypeNode *op) override;

        void VisitType_(const DictTypeNode *op) override;

        void VisitType_(const SetTypeNode *op) override;

        void VisitType_(const ExceptionTypeNode *op) override;

        void VisitType_(const IteratorTypeNode *op) override;

        void VisitType_(const FileTypeNode *op) override;

        void VisitType_(const NDArrayTypeNode *op) override;

        void VisitType_(const ClassTypeNode *op) override;

        void VisitType_(const UserDataTypeNode *op) override;

        void VisitType_(const OpaqueObjectTypeNode *op) override;

        void VisitType_(const RefTypeNode *op) override;
    };

    /*!
     * \brief TypeMutator that mutates expressions.
     */
    class TURBO_EXPORT TypeMutator : public TypeFunctor<Type(const Type &n)> {
    public:
        Type VisitType(const Type &t) override;

        Type VisitType_(const TypeVarNode *op) override;

        Type VisitType_(const FuncTypeNode *op) override;

        Type VisitType_(const TupleTypeNode *op) override;

        Type VisitType_(const GlobalTypeVarNode *op) override;

        Type VisitType_(const PrimTypeNode *op) override;

        Type VisitType_(const PointerTypeNode *op) override;

        Type VisitType_(const ObjectTypeNode *op) override;

        Type VisitType_(const UnicodeTypeNode *op) override;

        Type VisitType_(const StringTypeNode *op) override;

        Type VisitType_(const ListTypeNode *op) override;

        Type VisitType_(const DictTypeNode *op) override;

        Type VisitType_(const SetTypeNode *op) override;

        Type VisitType_(const ExceptionTypeNode *op) override;

        Type VisitType_(const IteratorTypeNode *op) override;

        Type VisitType_(const FileTypeNode *op) override;

        Type VisitType_(const NDArrayTypeNode *op) override;

        Type VisitType_(const ClassTypeNode *op) override;

        Type VisitType_(const UserDataTypeNode *op) override;

        Type VisitType_(const OpaqueObjectTypeNode *op) override;

        Type VisitType_(const RefTypeNode *op) override;

    private:
        runtime::Array<Type> MutateArray(runtime::Array<Type> arr);
    };

    /*!
     * \brief Bind free type variables in the type.
     * \param type The type to be updated.
     * \param args_map The binding map.
     */
    Type Bind(const Type &type, const runtime::Map<TypeVar, Type> &args_map);

}  // namespace tbir::ir
